import os

import numpy as np
import tensorflow as tf

"""
函数load_model：加载模型
checkpoint：.PB文件
"""

def load_model(checkpoint):
    """
    load the frozen graph of tensorflow as a detection model

    Parameters
    ----------
    checkpoint

    Returns
    -------

    """
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
    # 生成新的默认图
    detection_graph = tf.Graph()
    with detection_graph.as_default():
        # 新建GraphDef文件，用于临时载入模型中的图
        od_graph_def = tf.GraphDef()
        with tf.gfile.GFile(checkpoint, 'rb') as fid:
            # 二进制读取模型文件
            serialized_graph = fid.read()
            # GraphDef加载模型中的图
            od_graph_def.ParseFromString(serialized_graph)
            tf.import_graph_def(od_graph_def, name='')
    return detection_graph

"""
run_detection： 检测小图，返回box,class,score
sess：加载了模型的会话
detection_graph：加载了模型的图
image_np：小图
 
"""

def run_detection(sess, detection_graph, image_np):
    image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
    boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
    scores = detection_graph.get_tensor_by_name('detection_scores:0')
    classes = detection_graph.get_tensor_by_name('detection_classes:0')
    num_detections = detection_graph.get_tensor_by_name('num_detections:0')
    (boxes, scores, classes, num_detections) = sess.run(
        [boxes, scores, classes, num_detections],
        feed_dict={
            image_tensor: image_np
        })
    boxes = np.squeeze(boxes)
    classes = np.squeeze(classes).astype(np.int32)
    scores = np.squeeze(scores)
    return boxes, classes, scores
